﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.Remoting.Messaging;
using UnityEngine;
using UnityEngine.UI;

public class Main : MonoBehaviour
{
    public ComputeShader dofilterShader;
    public ComputeShader doDownSampleShader;
    public ComputeShader quanLianJieShader;

    public Shower theshow;

    public UnityEngine.UI.Slider slider;
    public Text t;

    private List<byte[]> rawImage;
    private List<byte[]> SizedImage;
    private List<byte[]> learnLabel;

    private List<byte[]> rewTestImage;
    private List<byte[]> SizedTestImage;
    private List<byte[]> TestLabel;


    //先卷积  计算完成后 再 池化(也就是 downsample  为 10 * 14 * 14 张图) 最后再全连接

    private float[,,,]     ws;                 //第一个纬度长度为10  第二个纬度 长度为 filterCount * 14 * 14    
    private float[]      bs;             //偏移值 长度为10    
    private float[,,,]     filters;        //卷积核 卷积核为6个  因此 长度为 10 * 6 * 9
    private float[,]     filterBs;       //卷积核的 偏移值     长度为 10 * 6

    //参数 总共为   10 * 10 * 14 * 14 + 10 * 10+ 10 * 6 * 9 + 10 * 6  

    public int filterCount = 2;
    public int imageSize = 28;
    public int filterWidth = 5;
    public int filterWidthHalf = 3;

    public const int LearnOnePatch = 30;    //一次性学习的图片数量

    private int CalcI = 0;          
    private int CalcJ = 0;
    private int CalcM = 0;

    float dw = 0.00001f;
    float df = 0.0001f;


    ComputeBuffer cbimages;
    ComputeBuffer cbw;
    ComputeBuffer cbwb;
    ComputeBuffer cbf ;
    ComputeBuffer cbfb;
    ComputeBuffer outvalues ;

    ComputeBuffer afterFliter;	//长度为 LearnOnePatch * FliterCount * 28 * 28
    ComputeBuffer afterdownsample;	//长度为 LearnOnePatch * FliterCount * 14 * 14

    // Use this for initialization
    void Start ()
    {
        Application.runInBackground = true;
        ws = new float[10, filterCount, 14 ,14];
        filters = new float[10,filterCount, 3 , 3];
        for (int i = 0; i < 10; i++)
        {
            for (int j = 0; j < filterCount; j++)
            {
                for (int k = 0; k < 3; k++)
                {
                    for (int l = 0; l < 3; l++)
                    {
                        filters[i, j, k, l] = 1.0f / 9.0f;
                    }
                }
            }
        }
        cbimages = new ComputeBuffer(LearnOnePatch * 30 * 30, 4);
        cbw = new ComputeBuffer(10 * filterCount * 14 * 14, 4);
        cbwb = new ComputeBuffer(10, 4); 
        cbf = new ComputeBuffer(10 * 9 * filterCount, 4);
        cbfb = new ComputeBuffer(10 * filterCount, 4);
        afterFliter = new ComputeBuffer(LearnOnePatch * filterCount * 10 * 28 * 28, 4);
        afterdownsample = new ComputeBuffer(LearnOnePatch * filterCount * 10 * 14 * 14, 4);
        outvalues = new ComputeBuffer(10 * LearnOnePatch, 4);             
        ReadW();
        LoadImage();
        theshow.InitSize(10,LearnOnePatch, 28, 3, filterCount, 14);
    }


    void OnDestroy()
    {
        cbimages.Dispose(); 
        cbw.Dispose(); 
        cbf.Dispose();
        cbfb.Dispose();
        cbwb.Dispose();
        afterFliter.Dispose();
        afterdownsample.Dispose();
        outvalues.Dispose();
    }

    // Update is called once per frame
    void Update ()
	{
	    slider.value = CalcI*0.1f + CalcJ*0.1f/14.0f;
	}

    public void LoadImage()
    {
        Reader r = new Reader();
        int width = 0;
        int height = 0;
        rawImage = r.LoadImage(r.trainimage, ref width, ref height);
        SizedImage = ReSize(rawImage, width, height);
        learnLabel = r.LoadLabel(r.trainlabel);

        rewTestImage = r.LoadImage(r.testimage, ref width, ref height);
        SizedTestImage = ReSize(rewTestImage, width, height);
        TestLabel = r.LoadLabel(r.testlabel);

        showW();
    }

    //将图片左右上下 各放大一个像素 并填上0
    List<byte[]> ReSize(List<byte[]> rawImage,int width,int height)
    {
        List < byte[] > images = new List<byte[]>();
        for (int i = 0; i < rawImage.Count; i++)
        {
            byte[] fixedImage = new byte[(width + 2)*(height + 2)];
            for (int j = 0; j < (width + 2); j++)
            {
                for (int k = 0; k < (height + 2); k++)
                {
                    if (j == 0 || k == 0 || j == width + 1 || k == height + 1)
                    {
                        fixedImage[j*(width + 2) + k] = 0;
                    }
                    else
                    {
                        fixedImage[j * (width + 2) + k] = rawImage[i][(j - 1) * width + (k - 1)];
                    }
                }
            }
            images.Add(fixedImage);
        }
        return images;
    }


    public void Learn()
    {
        StartCoroutine(LearnCo());
    }


    //得到 一个随机测试的图片， 并且保证 图片中包含10个数字
    int[] GetTestIndex()
    {
//        UnityEngine.Random.InitState(0);
        List<int> indexList = new List<int>();
        while (!isfullcontains(indexList.ToArray()))
        {
            indexList = new List<int>();
            //从图库中随机选择 N张图
            for (int i = 0; i < LearnOnePatch; i++)
            {
                //                indexList.Add(i);
                int aaaaa = UnityEngine.Random.Range(0, SizedImage.Count);
                while (indexList.Contains(aaaaa))
                {
                    aaaaa = UnityEngine.Random.Range(0, SizedImage.Count);
                }
                indexList.Add(aaaaa);
            }
        }
        int[] index = indexList.ToArray();
        return index;

    }

    bool isfullcontains(int[] index)
    {
        List<int> containsnumber = new List<int>();
        for(int i = 0 ; i < index.Length ; i++)
        {
            byte[] aa = learnLabel[index[i]];
            int k = -1;
            for(int j = 0 ; j < 10;j++)
            {
                if(aa[j] == 1)
                {
                    k = j;
                    break;
                }
            }
            if(!containsnumber.Contains(k))
            {
                containsnumber.Add(k);
            }            
        }
        if (containsnumber.Count == 10)
        {
            return true;
        }
        return false;
    }

    private bool isPause = false;

    private float[,] imagetemp;

    public IEnumerator LearnCo()
    {
        while (true)
        {
            showW();
            Test();

            int[] index = GetTestIndex();
            for (int i = 0; i < index.Length; i++)
            {
                byte[] aa = learnLabel[index[i]];
                int k = -1;
                for (int j = 0; j < 10; j++)
                {
                    if (aa[j] == 1)
                    {
                        Debug.Log(j);
                        break;
                    }
                }
            }
            imagetemp = new float[index.Length, SizedImage[0].Length];
            for (int i = 0; i < index.Length; i++)
            {
                for (int j = 0; j < SizedImage[index[i]].Length; j++)
                {
                    imagetemp[i, j] = SizedImage[index[i]][j];
                }
            }
            cbimages.SetData(imagetemp);

            theshow.SetRawImage(imagetemp, 30);

            double jiaochashang0 = GetTotalJiaoChaShang(ws, filters, cbimages, LearnOnePatch, index);

            Debug.Log("LearnOver oneTime");
            SaveW();
        }
    }

    public void Pause()
    {
        isPause = !isPause;
    }

    void showW()
    {

        /*
        float maxw = 0.0f;
        for (int i = 0; i < 10; i++)
        {
            for (int j = 0; j < 14; j++)
            {
                for (int m = 0; m < 14; m++)
                {
                    if (Math.Abs(ws[i , 0,j ,m]) > maxw)
                    {
                        maxw = Math.Abs(ws[i,0,j , m]);
                    }
                }
            }
        }
        float d = 1.0f / maxw * 0.95f;

        for (int i = 0; i < 10; i++)
        {
            wimage[i].texture = texturew[i];
            for (int j = 0; j < 14; j++)
            {
                for (int k = 0; k < 14; k++)
                {
                    float temp = ws[i,0,j , k];
                    Color c = new Color(temp < 0.0f ? (-temp * d) : 0.0f, temp > 0.0f ? (temp * d) : 0.0f, 0);
                    texturew[i].SetPixel(k,14 - j,c); 
                }
            }
            texturew[i].Apply();
        }
        */
    }

    void SaveW()
    {
        using(System.IO.BinaryWriter bw = new System.IO.BinaryWriter(System.IO.File.Open(Application.streamingAssetsPath + @"\w.bin",System.IO.FileMode.OpenOrCreate)))
        {
            Save4Array(bw,ws);
            SaveArray(bw, bs);
            Save4Array(bw, filters);
            Save2Array(bw, filterBs);
        }
    }

    void ReadW()
    {
        if(System.IO.File.Exists(Application.streamingAssetsPath + @"\w.bin"))
        {
            using(System.IO.BinaryReader br = new System.IO.BinaryReader(System.IO.File.Open(Application.streamingAssetsPath + @"\w.bin", System.IO.FileMode.OpenOrCreate)))
            {
                ws = Read4Array(br);
                bs = ReadArray(br);
                filters = Read4Array(br);
                filterBs = Read2Array(br);
            }
        }
        else
        {
            bs = new float[10];
            filters = new float[10, filterCount,3,3];
            for (int i = 0; i < filters.GetLength(0); i++)
            {
                for (int j = 0; j < filters.GetLength(1); j++)
                {
                    for (int k = 0; k < filters.GetLength(2); k++)
                    {
                        for (int l = 0; l < filters.GetLength(2); l++)
                        {
                            filters[i, j,k,l] = UnityEngine.Random.Range(-1.0f,1.0f);
                        }
                    }
                }
            }
            filterBs = new float[10, filterCount];

            ws = new float[10, filterCount , 14 , 14];
            for (int i = 0; i < ws.GetLength(0); i++)
            {
                for (int j = 0; j < ws.GetLength(1); j++)
                {
                    for (int k = 0; k < ws.GetLength(2); k++)
                    {
                        for (int l = 0; l < ws.GetLength(3); l++)
                        {
                            ws[i, j,k,l] = 0.0000001f;
                        }
                    }
                }
            }
        }

    }
    float[,] _lastEvidence;

    bool iswrong = false;

    double GetTotalJiaoChaShang(float[,,,] ws, float[,,,] fliters, ComputeBuffer images,int imageCount,int[] index)
    {
        if (iswrong)
        {
            return 0.0;
        }

        //首先将所有数据装配成显卡 用的格式
        cbw.SetData(ws);
        cbf.SetData(fliters);
        cbfb.SetData(filterBs);

        int kernel = dofilterShader.FindKernel("CSMain");
        dofilterShader.SetInt("ImageCount", imageCount);
        dofilterShader.SetInt("FliterCount", filterCount);
        dofilterShader.SetInt("FliterWidth", filterWidth);
        dofilterShader.SetInt("FliterWidthHalf", filterWidthHalf);
        dofilterShader.SetInt("ImageFixedSize", 14);

        dofilterShader.SetBuffer(kernel, "InImages", images);
        dofilterShader.SetBuffer(kernel, "InFliters", cbf);
        dofilterShader.SetBuffer(kernel, "InFlitersB", cbfb);

        dofilterShader.SetBuffer(kernel, "afterFliter", afterFliter);
        dofilterShader.Dispatch(kernel, imageCount, 10, 1);


        float[, , , ,] afterfliter = new float[imageCount, 10, filterCount, imageSize, imageSize];

        afterFliter.GetData(afterfliter);


        return 0.0f;
    }

    #region 相关float 数组的读取
    static float[,,,] Read4Array(System.IO.BinaryReader br)
    {
        int im = br.ReadInt32();
        int jm = br.ReadInt32();
        int km = br.ReadInt32();
        int lm = br.ReadInt32();
        float[,,,] outvalue = new float[im, jm, km, lm];
        for (int i = 0; i < outvalue.GetLength(0); i++)
        {
            for (int j = 0; j < outvalue.GetLength(1); j++)
            {
                for (int k = 0; k < outvalue.GetLength(2); k++)
                {
                    for (int l = 0; l < outvalue.GetLength(3); l++)
                    {
                        outvalue[i, j, k, l] = br.ReadSingle();
                    }
                }
            }
        }
        return outvalue;
    }
    static float[,,] Read3Array(System.IO.BinaryReader br)
    {
        int im = br.ReadInt32();
        int jm = br.ReadInt32();
        int km = br.ReadInt32();
        float[,,] outvalue = new float[im, jm, km];
        for (int i = 0; i < outvalue.GetLength(0); i++)
        {
            for (int j = 0; j < outvalue.GetLength(1); j++)
            {
                for (int k = 0; k < outvalue.GetLength(2); k++)
                {
                    outvalue[i, j, k] = br.ReadSingle();
                }
            }
        }
        return outvalue;
    }
    static float[,] Read2Array(System.IO.BinaryReader br)
    {
        int im = br.ReadInt32();
        int jm = br.ReadInt32();
        float[,] outvalue = new float[im, jm];
        for (int i = 0; i < outvalue.GetLength(0); i++)
        {
            for (int j = 0; j < outvalue.GetLength(1); j++)
            {
                outvalue[i, j] = br.ReadSingle();
            }
        }
        return outvalue;
    }
    static float[] ReadArray(System.IO.BinaryReader br)
    {
        int im = br.ReadInt32();
        float[] outvalue = new float[im];
        for (int i = 0; i < outvalue.GetLength(0); i++)
        {
            outvalue[i] = br.ReadSingle();
        }
        return outvalue;
    }

    static void Save4Array(System.IO.BinaryWriter bw, float[,,,] a)
    {
        bw.Write(a.GetLength(0));
        bw.Write(a.GetLength(1));
        bw.Write(a.GetLength(2));
        bw.Write(a.GetLength(3));
        for (int i = 0; i < a.GetLength(0);i++)
        {
            for (int j = 0; j < a.GetLength(1);j++)
            {
                for (int k = 0; k < a.GetLength(2);k++)
                {
                    for (int l = 0; l < a.GetLength(3);l++)
                    {
                        bw.Write(a[i,j,k,l]);
                    }
                }
            }
        }
    }
    static void Save3Array(System.IO.BinaryWriter bw, float[,,] a)
    {
        bw.Write(a.GetLength(0));
        bw.Write(a.GetLength(1));
        bw.Write(a.GetLength(2));
        for (int i = 0; i < a.GetLength(0); i++)
        {
            for (int j = 0; j < a.GetLength(1); j++)
            {
                for (int k = 0; k < a.GetLength(2); k++)
                {
                    bw.Write(a[i, j, k]);
                }
            }
        }
    }
    static void Save2Array(System.IO.BinaryWriter bw, float[,] a)
    {
        bw.Write(a.GetLength(0));
        bw.Write(a.GetLength(1));
        for (int i = 0; i < a.GetLength(0); i++)
        {
            for (int j = 0; j < a.GetLength(1); j++)
            {
                bw.Write(a[i, j]);
            }
        }
    }
    static void SaveArray(System.IO.BinaryWriter bw, float[] a)
    {
        bw.Write(a.GetLength(0));
        for (int i = 0; i < a.GetLength(0); i++)
        {
            bw.Write(a[i]);
        }
    }
    #endregion
    void Test()
    {
        List<int> wrongindex = new List<int>();

        int testcount = 1000;

        float[,] imagetemp = new float[testcount, SizedTestImage[0].Length];
        for (int i = 0; i < testcount; i++)
        {
            for (int j = 0; j < SizedTestImage[i].Length; j++)
            {
                imagetemp[i, j] = SizedTestImage[i][j];
            }
        }
        ComputeBuffer cbimagetemp = new ComputeBuffer(testcount * 30 * 30, 4);
        cbimagetemp.SetData(imagetemp);



        float[,] theE = TestAll(testcount, cbimagetemp);    //得到每个数字的权重
        cbimagetemp.Dispose();
        for (int i = 0 ; i < testcount; i++)
        {
            //计算出 可能性
            //计算softmax
            double totalsoftmax = 0.0f;
            for (int j = 0; j < 10; j++)
            {
                totalsoftmax += Math.Pow(Math.E, theE[j,i]);
            }
            double[] softmax = new double[10];
            for (int j = 0; j < 10; j++)
            {
                softmax[j] = Math.Pow(Math.E, theE[j,i]) / totalsoftmax;
            }


            int maxj = -1;
            float max = 0.0f;
            for (int j = 0; j < 10; j++)
            {
                if (softmax[j] > max)
                {
                    maxj = j;
                    max = (float)softmax[j];
                }
            }

            int maxlabel = -1;
            for (int j = 0; j < 10; j++)
            {
                if (TestLabel[i][j] == 1)
                {
                    maxlabel = j;
                    break;
                    ;
                }
            }

            if (maxlabel == maxj)
            {
                //预估正确;
            }
            else
            {
                wrongindex.Add(i);
            }
        }
        
        t.text = "错误率" + ((int)(1000 * wrongindex.Count / testcount)) / 10.0f ;
    }


    float[,] TestAll(int imageCount,ComputeBuffer cbimage)
    {
        ComputeBuffer outvaluesTest = new ComputeBuffer(10 * imageCount, 4);

        //首先将所有数据装配成显卡 用的格式
        cbw.SetData(ws);
        cbf.SetData(filters);
        cbwb.SetData(bs);
        cbfb.SetData(filterBs);

        int kernel = dofilterShader.FindKernel("CSMain");
        dofilterShader.SetInt("ImageCount", imageCount);
        dofilterShader.SetInt("FliterCount", filterCount);


        dofilterShader.SetBuffer(kernel, "InWs", cbw);
        dofilterShader.SetBuffer(kernel, "InBs", cbwb);

        dofilterShader.SetBuffer(kernel, "InImages", cbimage);
        dofilterShader.SetBuffer(kernel, "InFliters", cbf);
        dofilterShader.SetBuffer(kernel, "InFlitersB", cbfb);

        ComputeBuffer afterFlitertest = new ComputeBuffer(imageCount * filterCount * 10 * 28 * 28, 4);
        ComputeBuffer afterdownsampletest = new ComputeBuffer(imageCount * filterCount * 10 * 14 * 14, 4);


        dofilterShader.SetBuffer(kernel, "afterFliter", afterFlitertest);
        dofilterShader.SetBuffer(kernel, "afterdownsample", afterdownsampletest);

        dofilterShader.SetBuffer(kernel, "OutEvidence", outvaluesTest);

        dofilterShader.Dispatch(kernel, imageCount, 10, 1);

        float[,] outEvidence = new float[10, imageCount];
        //        float[,] outEvidence2 = new float[10, imageCount];

        outvaluesTest.GetData(outEvidence);
        outvaluesTest.Dispose();
        afterFlitertest.Dispose();
        afterdownsampletest.Dispose();
        return outEvidence;
    }


}
